#pip install xmltodict
import xmltodict
import json
#pip install pandasql
import pandas as pd
import re
from pandasql import sqldf
import xml.etree.ElementTree as ET
from ApiTools import apiTools, apiBase
#pip install 'polars[all]'
#pip install "unstructured[all-docs]"
# https://github.com/deanmalmgren/textract
#pip install textract ,pymupdf,pdfplumber


def map_csv(file_path):
    # 读取CSV文件
    df = pd.read_csv(file_path, sep=',')  # 如果是其他分隔符，比如制表符，可以将sep参数改为'\t'
    return df 


sql_reg1=r'if.*?end|else.*?end'
sql_reg2=r'select.*?from.*?where.*?;|delete.*?where.*?;|update.*?set.*?where.*?;|exec.*?;|create.*?;|drop.*?;'

#reg=r'if.*?end|else.*?end'
#reg=r'select.*?from.*?where.*?;|delete.*?where.*?;|update.*?set.*?where.*?;|exec.*?;'

def map_sql(file_path):
    sections_data = pd.DataFrame(columns=['Block', 'Content'])
    #file_path="./scripts/300.sql"
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
        content=content.lower()
        apiBase.lsExpress=[]
        content=apiBase.cut_all(content,apiBase.lsExpress,sql_reg1,sql_reg2,"----blockname")
        sections_data.append({'Block': "raw", 'Content': content})
        for i in range(len(apiBase.lsExpress)):            
            sections_data.append({'Block': f"block{i:05}", 'Content': apiBase.lsExpress[i]})
    
    # 创建DataFrame
    df = pd.DataFrame(sections_data)
    return df

#pip install pandas openpyxl
import pandas as pd
 
# 读取Excel文件
def map_excel(excel_file_path):
    df = pd.read_excel(excel_file_path)
    return df

# 假设你有一个XML字符串
# xml_data = """
# <root>
#     <element1>value1</element1>
#     <element2>value2</element2>
# </root>
# """

def map_xml(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        xml_data = file.read()    
        # 使用xmltodict.parse将XML转换为字典
        dict_data = xmltodict.parse(xml_data)
        
        # 使用json.dumps将字典转换为JSON格式的字符串
        json_data = json.dumps(dict_data, ensure_ascii=False, indent=4)
        # 读取JSON数据到DataFrame
        df = pd.read_json(json_data)        
        return df
 
# 假设您的JSON数据如下
# json_data = '''
# [
#     {"name": "Alice", "age": 25, "city": "New York"},
#     {"name": "Bob", "age": 30, "city": "San Francisco"}
# ]
# '''
def map_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        json_data = file.read()
        # 读取JSON数据到DataFrame
        df = pd.read_json(json_data)        
        return df
        
import pdfplumber
def map_pdf(file_path):
    # 创建一个空的DataFrame用于存储数据
    df = pd.DataFrame(columns=['Page', 'Content'])
    # 打开PDF文件
    with pdfplumber.open(file_path) as pdf:
        # 遍历每一页
        for page in pdf.pages:
            # 提取本页文本
            text = page.extract_text()
            df.loc[len(df)] = [page.page_number, text]                
    return df

#pip install python-docx
from docx import Document
def map_docx(docx_path):
    # 加载docx文件
    doc = Document(docx_path)
    # 初始化一个空列表来保存章节数据
    df = pd.DataFrame(columns=['line', 'Content'])
    paragraphs = doc.paragraphs
    index=0
    for paragraph in paragraphs:
        content = paragraph.text
        df.loc[len(df)] = [index, content]
        index=index+1
    return df

#pip install python-pptx
import pptx
from pptx import Presentation
# 定义一个函数来提取幻灯片文本
def extract_text_from_slide(slide):
    text_runs = [paragraph.text for shape in slide.shapes for paragraph in shape.text_frame.paragraphs]
    return '\n'.join(text_runs)

def  map_ppt(file_path): 
    # 加载PPT文件
    prs = Presentation(file_path)
    # 初始化一个列表来保存DataFrame
    df = pd.DataFrame(columns=['Page', 'Content'])
    # 遍历幻灯片并提取内容
    for i, slide in enumerate(prs.slides):
        text = extract_text_from_slide(slide)
        df.loc[len(df)] = [i, text]
    
    return df
# 创建一个简单的DataFrame
# df = pd.DataFrame({'name': ['Alice', 'Bob', 'Charlie'],
#                    'age': [25, 30, 35],
#                    'city': ['New York', 'Los Angeles', 'Chicago']})
file_path = apiBase.argv(1,"$PROJECT_HOME/200testdata/towngas.txt")
file_path=file_path.replace("'","")
#usr_prompt = apiBase.argv(2,"What brand of towngas do people ages between 30 and 40 buy?")
usr_prompt = apiBase.argv(2,"What brand of towngas do people aged between 30 and 40 buy?")

if file_path.endswith(".csv"):
    df = map_csv(file_path)
elif file_path.endswith(".txt"):
    df = map_csv(file_path)
elif file_path.endswith(".sql"):
    df = map_sql(file_path)
elif file_path.endswith(".ppt"): 
    map_ppt(file_path)
elif file_path.endswith(".pptx"): 
    df = map_ppt(file_path)
elif file_path.endswith(".docx"): 
    df = map_docx(file_path)
elif file_path.endswith(".doc"): 
    df = map_docx(file_path)
elif file_path.endswith(".pdf"):
    df = map_pdf(file_path)
elif file_path.endswith(".xml"):
    df = map_xml(file_path)
elif file_path.endswith(".json"):
    df = map_json(file_path)
elif file_path.endswith(".xlsx"):
    df = map_excel(file_path)
elif file_path.endswith(".xls"):
    df = map_excel(file_path)
else:
    print("Unsupported file format")
# 使用pandasql的sqldf函数执行SQL查询
pysqldf = lambda q: sqldf(q, globals())
#usr_prompt = "What brand of towngas do people aged between 30 and 40 buy?"
try:
    names=df.columns.tolist()
    column_names = ', '.join(names)
    #name ,age ,brand name后面不能有空格
    #result = pysqldf("select * from df where age > 10")
    if len(usr_prompt) == 0:
        query = f"SELECT {column_names} FROM df "        
    else:        
        sys_prompt = f'''You are a SQL expert with a strong attention to detail.
1. Double check the SQLite query for common mistakes.
2. The table name is: df .
3. The fields is: {column_names} .

Generate SQL statements based on table name and fields'''
        prompt = f'''## content
{usr_prompt}
        
## requirements
According to the above content, convert to SQL with the following requirements:
1. Please convert the codes line by line, do not skip any line of code.
2. The output should only provide the PostgreSQL codes without any explanations.
3. The output should execute without any errors.
4. Do not change any variable names.
5. The correct code output format is
```code```'''
        query=apiTools.complete(sys_prompt,prompt)
        #query = "SELECT age,brand FROM df WHERE age > 20;"
        query=query.lower()
        structures = re.findall(r'select.*?```', query, re.DOTALL)
        # 输出截取的表结构字符串
        for st in structures:
            query=st[0:-3]
    
    result = pysqldf(query)
    column_names = ', '.join(result.columns)
    print(column_names)
    
    for index, row in result.iterrows():        
        itemval=''
        for name in result.columns:
            itemval=itemval+','+str(row[name])
        itemval=itemval[1:]+"\n"
        print(itemval)
finally:
    apiBase.close()
